#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
WebSocket客户端测试脚本
用于测试WebSocket代理桥接功能
"""

import asyncio
import websockets
import json
import logging
import argparse
from datetime import datetime

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class WebSocketTestClient:
    """WebSocket测试客户端"""
    
    def __init__(self, server_url: str, proxy_host: str = None, proxy_port: int = None):
        """
        初始化WebSocket测试客户端
        
        Args:
            server_url: WebSocket服务器URL
            proxy_host: 代理服务器地址
            proxy_port: 代理服务器端口
        """
        self.server_url = server_url
        self.proxy_host = proxy_host
        self.proxy_port = proxy_port
        self.websocket = None
        self.running = False
    
    async def connect(self):
        """连接到WebSocket服务器"""
        try:
            # 如果指定了代理，则通过代理连接
            if self.proxy_host and self.proxy_port:
                # 注意：websockets库不直接支持HTTP代理
                # 这里只是示例，实际使用时可能需要其他方法
                logger.info(f"尝试通过代理 {self.proxy_host}:{self.proxy_port} 连接到 {self.server_url}")
            else:
                logger.info(f"直接连接到 {self.server_url}")
            
            self.websocket = await websockets.connect(self.server_url)
            self.running = True
            logger.info("WebSocket连接已建立")
            
        except Exception as e:
            logger.error(f"连接失败: {str(e)}")
            raise
    
    async def send_message(self, message: str):
        """发送消息"""
        if self.websocket and not self.websocket.closed:
            await self.websocket.send(message)
            logger.info(f"已发送消息: {message}")
        else:
            logger.error("WebSocket连接未建立或已关闭")
    
    async def listen_messages(self):
        """监听接收消息"""
        try:
            while self.running and self.websocket:
                message = await self.websocket.recv()
                timestamp = datetime.now().strftime('%H:%M:%S')
                logger.info(f"[{timestamp}] 收到消息: {message}")
                
        except websockets.exceptions.ConnectionClosed:
            logger.info("WebSocket连接已关闭")
        except Exception as e:
            logger.error(f"接收消息时出错: {str(e)}")
    
    async def send_periodic_messages(self, interval: int = 5):
        """定期发送测试消息"""
        counter = 1
        while self.running:
            try:
                test_message = json.dumps({
                    "type": "test",
                    "counter": counter,
                    "timestamp": datetime.now().isoformat(),
                    "message": f"这是第 {counter} 条测试消息"
                }, ensure_ascii=False)
                
                await self.send_message(test_message)
                counter += 1
                await asyncio.sleep(interval)
                
            except Exception as e:
                logger.error(f"发送定期消息时出错: {str(e)}")
                break
    
    async def run_interactive_mode(self):
        """运行交互模式"""
        logger.info("进入交互模式，输入消息后按回车发送，输入 'quit' 退出")
        
        while self.running:
            try:
                # 在实际应用中，这里需要使用异步输入
                # 这里简化处理，实际可以使用 aioconsole 库
                await asyncio.sleep(1)
                
            except KeyboardInterrupt:
                logger.info("收到中断信号，退出交互模式")
                break
    
    async def close(self):
        """关闭连接"""
        self.running = False
        if self.websocket:
            await self.websocket.close()
            logger.info("WebSocket连接已关闭")


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='WebSocket客户端测试工具')
    parser.add_argument('--url', default='ws://echo.websocket.org/', 
                       help='WebSocket服务器URL (默认: ws://echo.websocket.org/)')
    parser.add_argument('--proxy-host', help='代理服务器地址')
    parser.add_argument('--proxy-port', type=int, help='代理服务器端口')
    parser.add_argument('--mode', choices=['interactive', 'auto'], default='auto',
                       help='运行模式: interactive(交互) 或 auto(自动测试)')
    parser.add_argument('--interval', type=int, default=5,
                       help='自动模式下发送消息的间隔秒数')
    
    args = parser.parse_args()
    
    # 创建客户端
    client = WebSocketTestClient(args.url, args.proxy_host, args.proxy_port)
    
    try:
        # 连接到服务器
        await client.connect()
        
        # 启动消息监听任务
        listen_task = asyncio.create_task(client.listen_messages())
        
        if args.mode == 'interactive':
            # 交互模式
            await client.run_interactive_mode()
        else:
            # 自动测试模式
            logger.info(f"启动自动测试模式，每 {args.interval} 秒发送一条消息")
            await client.send_periodic_messages(args.interval)
        
    except KeyboardInterrupt:
        logger.info("收到中断信号")
    except Exception as e:
        logger.error(f"运行时出错: {str(e)}")
    finally:
        await client.close()


if __name__ == "__main__":
    asyncio.run(main())